package com.ycky.batchProcess;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @description: 发布速率限制器
 * @author: cjj
 * @time: 2025/3/18 10:59
 */
public class GlobalRateLimiter {
    private final AtomicLong lastRequestTime = new AtomicLong(System.currentTimeMillis());
    private final AtomicInteger currentSecondCount = new AtomicInteger(0);
    private final AtomicLong currentSecond = new AtomicLong(System.currentTimeMillis() / 1000);

    public void acquire(int rate) throws InterruptedException {
        long now = System.currentTimeMillis();
        long currentSec = now / 1000;

        // 如果进入新的一秒，重置计数器
        if (currentSec > currentSecond.get()) {
            currentSecond.set(currentSec);
            currentSecondCount.set(0);
        }

        // 检查是否超过速率限制
        int count = currentSecondCount.incrementAndGet();
        if (count > rate) {
            // 计算需要等待的时间
            long nextSecond = (currentSec + 1) * 1000;
            long waitTime = nextSecond - now;
            if (waitTime > 0) {
                Thread.sleep(waitTime);
            }
            // 重置计数器
            currentSecond.set(System.currentTimeMillis() / 1000);
            currentSecondCount.set(1);
        } else {
            // 计算请求间隔
            long interval = 1000 / rate;
            long elapsed = now - lastRequestTime.get();
            if (elapsed < interval) {
                Thread.sleep(interval - elapsed);
            }
        }

        lastRequestTime.set(System.currentTimeMillis());
    }
}